package cn.ycc1.boot3djl.controller;

import lombok.SneakyThrows;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;

import ai.djl.MalformedModelException;
import ai.djl.Model;
import ai.djl.basicdataset.cv.classification.ImageFolder;
import ai.djl.basicdataset.cv.classification.Mnist;
import ai.djl.basicmodelzoo.basic.Mlp;
import ai.djl.inference.Predictor;
import ai.djl.metric.Metrics;
import ai.djl.modality.Classifications;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.ImageFactory;
import ai.djl.modality.cv.transform.CenterCrop;
import ai.djl.modality.cv.transform.Resize;
import ai.djl.modality.cv.transform.ToTensor;
import ai.djl.modality.cv.translator.ImageClassificationTranslator;
import ai.djl.ndarray.NDArray;
import ai.djl.ndarray.NDManager;
import ai.djl.ndarray.types.Shape;
import ai.djl.training.DefaultTrainingConfig;
import ai.djl.training.EasyTrain;
import ai.djl.training.Trainer;
import ai.djl.training.TrainingResult;
import ai.djl.training.dataset.Dataset;
import ai.djl.training.dataset.RandomAccessDataset;
import ai.djl.training.evaluator.Accuracy;
import ai.djl.training.listener.TrainingListener;
import ai.djl.training.loss.Loss;
import ai.djl.training.util.ProgressBar;
import ai.djl.translate.TranslateException;

import java.io.IOException;
import java.nio.file.Path;
import java.nio.file.Paths;


/**
 * 测试djl
 *
 * @author ycc
 * @date 2024/12/9
 *
 * MLP 的工作原理可以分为前向传播（Forward Propagation）和反向传播（Backward Propagation）
 * 两个主要步骤：
 * 前向传播：
 *   ○ 输入数据从输入层传递到输出层。
 *   ○ 每个神经元计算加权和（线性组合），然后通过激活函数引入非线性。
 * 反向传播：
 *   ○ 反向传播用于计算损失函数相对于每个权重和偏置的梯度。
 *   ○ 使用梯度下降法更新权重和偏置，以最小化损失函数。
 * 具体步骤包括：
 * ● 计算输出层的误差。
 * ● 逐层向后传播误差，计算每个隐藏层的误差。
 * ● 更新权重和偏置。
 *
 * 当我们准备好数据集和神经网络之后，就可以开始训练模型了。在深度学习中，一般会由下面几步来完成一个训练过程：
 * ● 初始化：我们会对每一个 Block 的参数进行初始化，初始化每个参数的函数都是由 设定的 Initializer 决定的。
 * ● 前向传播：这一步将输入数据在神经网络中逐层传递，然后产生输出数据。
 * ● 计算损失：我们会根据特定的损失函数 Loss 来计算输出和标记结果的偏差。
 * ● 反向传播：在这一步中，你可以利用损失反向求导算出每一个参数的梯度。
 * ● 更新权重：我们会根据选择的优化器（Optimizer）更新每一个在 Block 上参数的值。
 *
 * Dataset - 数据集
 * 数据集是用于训练机器学习模型的数据集合。
 * 机器学习通常使用三个数据集：https://machinelearningmastery.com/difference-test-validation-datasets/
 * ● 训练数据集
 *   ○ 我们用来训练模型的实际数据集。模型从这些数据中学习权重和参数。
 * ● 验证数据集
 *   ○ 验证集用于在训练过程中评估给定模型。它帮助机器学习工程师在模型开发阶段微调超参数。模型不从验证数据集学习；验证数据集是可选的。
 * ● 测试数据集
 *   ○ 测试数据集提供了用于评估模型的黄金标准，它只在模型完全训练后使用，测试数据集应该更准确地评估模型将如何在新数据上执行。
 *
 *  https://docs.djl.ai/master/docs/dataset.html
 *
 * 由于鞋分类模型训练时间会比较长，建议使用显卡，需要安装 CUDA、cuDNN环境
 * nvcc -V确认cuda版本
 *
 *  安装pytorch环境
 * pip3 install torch torchvision torchaudio --user --index-url https://download.pytorch.org/whl/cu121
 *
 * CUDA
 * https://developer.nvidia.com/cuda-toolkit
 * https://developer.nvidia.com/rdp/cudnn-download
 * cuda：下载安装即可
 * cudnn：下载后，解压把 bin、lib等目录复制到cuda的安装目录下
 */
@RestController
public class TestController {

    // 测试模型
    @GetMapping("/predict")
    void predictTest() throws IOException, MalformedModelException, TranslateException {
        // 1、准备测试数据
        Image image = ImageFactory.getInstance().fromUrl("https://resources.djl.ai/images/0.png");
        Image img2 = ImageFactory.getInstance().fromFile(Paths.get("build/img/5.png"));

        // 2、加载模型
        Path modelDir = Paths.get("build/mlp");
        Model model = Model.newInstance("mlp");
        model.setBlock(new Mlp(28 * 28, 10, new int[] {128, 64}));
        model.load(modelDir);


        // 3、获取一个转化器； 黑底白字， 白底黑字
        ImageClassificationTranslator translator = ImageClassificationTranslator.builder()
                .addTransform(new Resize(28, 28))
                //加处理逻辑
                .addTransform(new ToTensor())
                .build();

        // 4、获取预测器
        Predictor<Image, Classifications> predictor = model.newPredictor(translator);

        // 5、预测图片分类
        Classifications predict = predictor.predict(img2);

        System.out.println(predict);
    }

    // 完全训练一个模型：
    // 1、准备数据集
    // 2、构建神经网络
    // 3、构建模型（这个模型应用上面的神经网络）
    // 4、训练配置（如何训练、训练集、验证集、测试集）
    // 5、保存模型
    // ======下面流程===
    // 6、加载模型
    // 7、预测（给模型一个新输入，让他判断这是什么）
    @SneakyThrows
    @GetMapping("/training")
    void trainingComplete() throws Exception {
        // 1、准备数据集; 用自己的数据就自定义 Dataset
        RandomAccessDataset trainingSet = getDataset(Dataset.Usage.TRAIN);
        RandomAccessDataset validateSet = getDataset(Dataset.Usage.TEST);

        // 2、构建神经网络； 就是 block
        // 输入：28*28  输出：10  隐藏层：128 64
        Mlp mlp = new Mlp(Mnist.IMAGE_WIDTH * Mnist.IMAGE_WIDTH, Mnist.NUM_CLASSES, new int[]{128, 64});

        // 3、创建模型，加载这个block
        try (Model model = Model.newInstance("mlp")){
            // 设置神经网络
            model.setBlock(mlp);

            // 4、训练配置
            String outputDir = "build/mlp";
            DefaultTrainingConfig config = new DefaultTrainingConfig(Loss.softmaxCrossEntropyLoss())
                    .addEvaluator(new Accuracy())
                    .addTrainingListeners(TrainingListener.Defaults.logging(outputDir));

            // 5、获取训练器
            try (Trainer trainer = model.newTrainer(config)) {
                trainer.setMetrics(new Metrics());
                // 6、初始化
                trainer.initialize(new Shape(1, Mnist.IMAGE_HEIGHT * Mnist.IMAGE_WIDTH));

                // 7、训练5次
                EasyTrain.fit(trainer, 5, trainingSet, validateSet);

                TrainingResult result = trainer.getTrainingResult();
                System.out.println("训练结果..." + result);

                // 8、保存模型  DJL
                model.save(Paths.get(outputDir),"mlp");

                System.out.println("模型保存成功");
            }
        };
    }

    @SneakyThrows
    private RandomAccessDataset getDataset(Dataset.Usage usage) throws Exception{
        //1、Mnist 是内置的，数据已经带好了 用户目录/.djl/cache

        //这个就是数据集、自定义的数据集，需要写translator
//        ImageFolder build = ImageFolder.builder()
//                .setRepositoryPath("~/.djl/cache")
//                .optImageHeight(28)
//                .optImageWidth(28)
//                .setSampling(64,true)
//                .build();
//
        Mnist mnist = Mnist.builder()
                .setSampling(64, true)
                .optUsage(usage)
                .build();
        mnist.prepare(new ProgressBar());
        return mnist;
    }

}
